import argparse
import torch

parser = argparse.ArgumentParser(description='GC-Net')
parser.add_argument('--dataset', default='kitti2015', help='datapath', choices=['kitti2015', 'sceneflow', 'usvinland', 'usvinland_seg'])
parser.add_argument('--batch_size', type=int, default=1, help='batch_size')
parser.add_argument('--num_workers', type=int, default=0, help='num_workers')
parser.add_argument('--epochs', type=int, default=50, help='number of epochs to train')
args = parser.parse_args()

# 操作系统 Windows, Ubuntu, WSL
system = 'Windows'
if system == 'Windows': sys_root = 'G:'
elif system == 'Ubuntu': sys_root = '/media/ubuntu/e/zhouyiqing'
elif system == 'WSL': sys_root = '/mnt/g'

# 数据集 kitti2015, sceneflow, usvinland
dataset = args.dataset
seg = False
if dataset == 'usvinland_seg':
    dataset = 'usvinland'
    seg = True

# 是否加载预训练权重
load_pretrained = True

# 是否随机裁切
random_crop = False

# 是否下采样
under_sampling = False

# 训练参数
BATCH_SIZE = args.batch_size
NUM_WORKERS = args.num_workers # Window上为0
NUM_EPOCHS = args.epochs
device = torch.device('cuda')

# 控制输入尺寸(宽高应为32的倍数)
if random_crop:
    # 随机裁切(256, 512) max_disp = 192 (第三方 96)
    # Ubuntu尝试Scene Flow(480, 768) max_disp = 256
    height = 256
    width = 512
    max_disp = 192
else:
    # 全尺寸：kitti2015(375, 1242), sceneflow(540, 960), usvinland(320, 640)
    # kitti2015 = 229(192), sceneflow = 256(192), usvinland = 64
    if dataset == 'kitti2015':
        height = 352
        width = 1216
        max_disp = 192
    elif dataset == 'sceneflow':
        height = 512
        width = 960
        max_disp = 192
    elif dataset == 'usvinland':
        height = 320
        width = 640
        max_disp = 64

    # 全尺寸2倍下采样
    if under_sampling:
        height = int(height /2)
        width = int(width / 2)
        max_disp = int(max_disp / 2)
